import os
import shutil

import imagehash
from PIL import Image
import numpy as np
import multiprocessing as mp
from itertools import combinations

"""
输入参数：
1、图片的文件夹加载路径
2、重复图片的保存路径
3、阈值
输出参数：
1、列表（列表内包含重复图片的全路径）
2、去除重复图片之后的路径
3、重复图片的保存路径

"""
counter = mp.Value('i', 0)


def map(samples):
    a = []
    global counter
    for sample in samples:
        c_hash = imagehash.phash(Image.open(sample))
        # c_hash = imagehash.phash(Image.open(sample)).convert('RGB')
        a.append([c_hash, sample])
        with counter.get_lock():
            counter.value += 1
            # if counter.value % 10 == 0:
            # print(counter.value)
    return a


def gen_hash(img_data_path):
    # 获取每张图片的hash值
    image_names = os.listdir(img_data_path)
    image_names = np.sort(image_names)
    all_img_paths = [os.path.join(img_data_path, x) for x in image_names]
    # split_names = np.split(all_img_paths, [100, 200, 300, 400, 500, 600, 700, 800, 900])

    n_thread = 1

    if n_thread == 1:
        split_names = [all_img_paths]
        # p／rint("split_names".split_names)
    else:
        n = len(all_img_paths)
        jiange = int(n / n_thread)
        split_point = []
        for i in range(1, n_thread):
            split_point.append(i * jiange)
        split_names = np.split(all_img_paths, split_point)
    processS = mp.Pool(processes=len(split_names))

    c_results = [[] for i in range(len(split_names))]

    queue = mp.Manager().Queue()
    # 枚举函数；i表示循环下标；block表示从迭代器获取到的元素
    for i, block in enumerate(split_names):
        c_results[i] = processS.apply_async(func=map, args=(block,))
    processS.close()
    processS.join()

    s = []

    for i, values in enumerate(c_results):
        values = values.get()
        for value in values:
            s.append('{} {}'.format(value[0], value[1]))

    return s


def dup_img(sim_thres, hash):
    # 根据hash值查找重复的图片
    # result = {}
    sim_thres = 0 if sim_thres < 0 else sim_thres

    # if output_path == '':
    #     output_path = os.path.join(os.path.split(pHashs_path)[0], 'duplicated_images.json')
    s = []
    # f.readline()
    for x in hash:
        s.append(x.strip().split(' '))

    name_phash_mapping = {}
    for x in s:
        name_phash_mapping[x[1]] = '{:064d}'.format(int(bin(int(x[0], 16))[2:]))

    n_split = 4
    # Error
    # dicts = [{}] * n_split
    dicts = [{} for i in range(n_split)]

    for i in range(n_split):
        for x in s:
            c_part = x[0][i * int(16 / n_split):(i + 1) * int(16 / n_split)]
            if c_part not in dicts[i]:
                dicts[i][c_part] = []
            assert x[1] not in dicts[i][c_part], x[1]
            dicts[i][c_part].append(x[1])
    # print(dicts)

    repeated_img_urls = {}

    count = 0
    for i in range(n_split):
        for img_list in dicts[i].values():
            count += 1
            # if count % 100 == 0:
            #     print(count)

            if len(img_list) < 2: continue

            for img1, img2 in combinations(img_list, 2):
                assert img1 != img2, '{}\n{}'.format(img1, img2)
                img1_finger = name_phash_mapping[img1]
                img2_finger = name_phash_mapping[img2]
                num_diff = sum([img1_finger[i] != img2_finger[i] for i in range(64)])
                # print(num_diff)
                if num_diff <= int(sim_thres):  # The number of different finger digits
                    if img1 not in repeated_img_urls and img2 not in repeated_img_urls:
                        t_img_set = set([img1, img2])
                        repeated_img_urls[img1] = t_img_set
                        repeated_img_urls[img2] = t_img_set  # The same reference
                    elif img1 not in repeated_img_urls:
                        repeated_img_urls[img2].add(img1)
                        repeated_img_urls[img1] = repeated_img_urls[img2]
                    elif img2 not in repeated_img_urls:
                        repeated_img_urls[img1].add(img2)
                        repeated_img_urls[img2] = repeated_img_urls[img1]
                    else:
                        repeated_img_urls[img1].add(img2)
                        repeated_img_urls[img2].add(img1)
                        repeated_img_urls[img1] = repeated_img_urls[img1] | repeated_img_urls[img2]

                        # Any other keys will change to this set.
                        # Broadcast to other imgs
                        for img in repeated_img_urls[img1]:
                            repeated_img_urls[img] = repeated_img_urls[img1]

    for x in repeated_img_urls.values():
        listx = list(x)
        assert len(listx) == len(x)
        c = repeated_img_urls[listx[0]]
        for y in listx:
            assert repeated_img_urls[y] == c

    repeated_img_urls_set = []
    for x in repeated_img_urls:
        assert len(repeated_img_urls[x]) >= 2
        if repeated_img_urls[x] not in repeated_img_urls_set:
            repeated_img_urls_set.append(repeated_img_urls[x])

    repeated_img_urls_list = []
    for x in repeated_img_urls_set:
        repeated_img_urls_list.append(list(x))

    sum1 = sum([len(x) for x in repeated_img_urls_list])

    a = set()
    for x in repeated_img_urls_list:
        for y in x:
            a.add(y)

    assert len(a) == sum1, 'Unknown error.'
    # print(repeated_img_urls_list)
    return repeated_img_urls_list
    # import json
    # json.dump(repeated_img_urls_list, open(output_path, 'w'))


def phashdup(save_folder,load_folder,threshold):
    os.makedirs(save_folder, exist_ok=True)
    try:
        result = {}
        img_data_path = load_folder
        imgs_nums = len(os.listdir(img_data_path))

        # n_thread = arg["n_thread"]
        # out_path = arg["out_path"]
        hash = gen_hash(img_data_path)
        # print(hash)
        sim_thres = threshold
        # output_path = arg["out_path"]
        # pHashs_path = arg["pHashs_path"]
        dup_imgs = dup_img(sim_thres, hash)
        # print(dup_imgs)
        result["msg"] = "检测成功"
        result["code"] = 0
        result["algo"] = "phashdup"
        # result["imgs_nums"] = imgs_nums
        # result["sim_nums"] = len(dup_imgs)
        for dup, i in zip(dup_imgs, range(len(dup_imgs))):
            # print(i)
            result["sim", i] = dup
            for j in range(1,len(dup)):
                shutil.move(dup[j], save_folder) # 保留第一张,其余的移动到新的文件夹里面
            # print(dup)

        return result
    except:
        result["msg"] = "检测失败"
        result["code"] = -1
        return result


